package org.syntaxlisp.core.match.tree;

import lombok.Builder;
import lombok.Data;
import lombok.Getter;
import org.apache.commons.collections4.CollectionUtils;
import org.syntaxlisp.core.ParseException;
import org.syntaxlisp.core.element.*;
import org.syntaxlisp.core.match.MatchUnit;
import org.syntaxlisp.core.utils.Pair;

import java.util.*;
import java.util.stream.Collectors;

import static org.syntaxlisp.core.match.MatchUnitType.STRING;

/**
 * @author HuangYijun
 * @data 2023-07-07 10:42
 **/
@Data
public class MatchState {

    private DefID defID;

    private DefID belongsToDef;

    private SyntaxLispElement src;

    private boolean canSkipToUpperNext;

    private boolean hunger;

    private boolean matchButIgnore;

    private Set<SyntaxLispElement> lookAheadExcludes;

    StringTreeMap<Boolean> lookAheadExcludesMap;

    private boolean reduceAsString;

    private MatchState upper;

    private Map<MatchState, MatchState> baseNextBelongsToDirectNextMap;

    private Set<MatchState> hungerStateSet;

    /**
     * base match map 记录当前 state base match map
     * 当 state 不为终止 state 或 src element 为 REF 类型时，base match map 不为空
     */
    Map<DefID, Map<MatchState, Set<MatchState>>> pfBaseJumpMap;

    StringTreeMap<Map<MatchState, Set<MatchState>>> pfStringJumpMap;

    Map<Character, Map<MatchState, Set<MatchState>>> pfCharacterJumpMap;

    /**
     * 线性表中的键值对数据为 匹配目标状态集
     */
    // todo: baseJumpMap key change
    Map<DefID, Map<MatchState, Set<MatchState>>> nextBaseJumpMap;

    StringTreeMap<Map<MatchState, Set<MatchState>>> nextStringJumpMap;

    Map<Character, Map<MatchState, Set<MatchState>>> nextCharacterJumpMap;

    public void initMap (Map<SyntaxLispElement, MatchState> elementStateMap, Map<String, DefID> defIdMap) {
        Objects.requireNonNull(src, "Source element can not be null before init map");
        pfBaseJumpMap = new HashMap<>();
        pfCharacterJumpMap = new HashMap<>();
        pfStringJumpMap = new StringTreeMap<>();
        nextBaseJumpMap = new HashMap<>();
        nextCharacterJumpMap = new HashMap<>();
        nextStringJumpMap = new StringTreeMap<>();
        baseNextBelongsToDirectNextMap = new HashMap<>();
        hungerStateSet = new HashSet<>();
        genMapFromPriorityList (src.getBasePriorityPfList(),
                                pfBaseJumpMap, pfCharacterJumpMap, pfStringJumpMap,
                                elementStateMap, defIdMap);
        genMapFromPriorityList (src.getBasePriorityNextList(),
                                nextBaseJumpMap, nextCharacterJumpMap, nextStringJumpMap,
                                elementStateMap, defIdMap);
        if (src.isRef() || src.isBase()) {
            defID = defIdMap.get(src.getVal());
        }
    }

    // todo: 2024/6/30 简化修改
    public static void genMapFromPriorityList (List<Map<SyntaxLispElement, Set<SyntaxLispElement>>> priorityList,
                                               Map<DefID, Map<MatchState, Set<MatchState>>> srcBaseMap,
                                               Map<Character, Map<MatchState, Set<MatchState>>> srcCharacterMap,
                                               Map<String, Map<MatchState, Set<MatchState>>> srcStringMap,
                                               Map<SyntaxLispElement, MatchState> elementStateMap,
                                               Map<String, DefID> defIDMap) {
        if (!CollectionUtils.isEmpty(priorityList)) {
            priorityList.forEach(curPriorityPair -> {
                Map<DefID, Map<MatchState, Set<MatchState>>> baseMap = new HashMap<>();
                Map<Character, Map<MatchState, Set<MatchState>>> characterMap = new HashMap<>();
                Map<String, Map<MatchState, Set<MatchState>>> stringMap = new HashMap<>();
                curPriorityPair.forEach((matchElement, localElementSet) -> {
                    MatchState targetState = elementStateMap.get(matchElement);
                    Set<MatchState> localStateSet = localElementSet.stream()
                                                                   .map(elementStateMap::get)
                                                                   .collect(Collectors.toSet());
                    switch (matchElement.getType()) {
                        case BASE: {
                            DefID baseID = defIDMap.get(matchElement.getVal());
                            baseMap.computeIfAbsent(baseID, k -> new HashMap<>())
                                   .computeIfAbsent(targetState, k -> new HashSet<>())
                                   .addAll(localStateSet);
                            break;
                        }
                        case CHARACTER: {
                            characterMap.computeIfAbsent(((CharacterElement) matchElement).getCharVal(),
                                                         k -> new HashMap<>())
                                        .computeIfAbsent(targetState, k -> new HashSet<>())
                                        .addAll(localStateSet);
                            break;
                        }
                        case SCOPE: {
                            for (Integer codepoint : ((MapScopeElement) matchElement).getAllCodepointInScope()) {
                                characterMap.computeIfAbsent((char) codepoint.intValue(), k -> new HashMap<>())
                                            .computeIfAbsent(targetState, k -> new HashSet<>())
                                            .addAll(localStateSet);
                            }
                            break;
                        }
                        case STRING: {
                            stringMap.computeIfAbsent(matchElement.getVal(), k -> new HashMap<>())
                                     .computeIfAbsent(targetState, k -> new HashSet<>())
                                     .addAll(localStateSet);
                            break;
                        }
                        default: {
                            throw new ParseException("Un-support element type in TreeState initMap operation.");
                        }
                    }
                });
                putToStateMapIfAbsent(srcBaseMap, baseMap);
                putToStateMapIfAbsent(srcCharacterMap, characterMap);
                putToStateMapIfAbsent(srcStringMap, stringMap);
            });
        }
    }

    public static <K> void putToStateMapIfAbsent(Map<K, Map<MatchState, Set<MatchState>>> srcMap,
                                                 Map<K, Map<MatchState, Set<MatchState>>> targetMap) {
        targetMap.forEach((key, stateSet) -> {
            if (!srcMap.containsKey(key)) {
                srcMap.put(key, new HashMap<>(stateSet));
            }
        });
    }

    public
    List<Pair<Map<MatchState, Set<MatchState>>, Integer>>
    matchPf(List<MatchUnit> inputSeq,
            int offset,
            Map<DefID, Map<MatchState, Set<MatchState>>> baseJumpMap,
            StringTreeMap<Map<MatchState, Set<MatchState>>> stringJumpMap,
            Map<Character, Map<MatchState, Set<MatchState>>> characterJumpMap,
            Map<Integer, Set<DefID>> predicateMap,
            boolean usePredicateMap) {

        if (offset >= inputSeq.size()) {
            return Collections.emptyList();
        }

        List<Pair<Map<MatchState, Set<MatchState>>, Integer>> res = new ArrayList<>(2);
        MatchUnit firstUnit = inputSeq.get(offset);
        switch (firstUnit.getType()) {
            case BASE: {
                Map<MatchState, Set<MatchState>> jumpMap = baseJumpMap.getOrDefault(firstUnit.getBaseId(),
                                                                                    Collections.emptyMap());
                if (usePredicateMap) {
                    jumpMap = filterJumpMap(offset, jumpMap, predicateMap);
                }
                if (!jumpMap.isEmpty()) {
                    res.add(Pair.<Map<MatchState, Set<MatchState>>, Integer>builder()
                                      .first(jumpMap)
                                      .second(offset + 1)
                                      .build());
                }
                break;
            }
            case STRING: {
                // 支持多个 inputSeqUnit 匹配单个 state（element），不支持多个 state（element）匹配单个 inputSeqUnit
                List<Pair<String, Integer>> matchKeyResList = stringJumpMap.matchKey(inputSeq, offset);
                for (Pair<String, Integer> matchKeyRes : matchKeyResList) {
                    if (matchKeyRes != null
                            && offset < matchKeyRes.getSecond()) {
                        String matchString = matchKeyRes.getFirst();
                        Map<MatchState, Set<MatchState>> jumpMap = stringJumpMap.get(matchString);
                        if (usePredicateMap) {
                            jumpMap = filterJumpMap(offset, jumpMap, predicateMap);
                        }
                        res.add(Pair.<Map<MatchState, Set<MatchState>>, Integer>builder()
                                          .first(jumpMap)
                                          .second(matchKeyRes.getSecond())
                                          .build());
                    }
                }
                // 尝试匹配 characterJumpMap
                Character val = tryGetCharVal(firstUnit);
                if (val != null) {
                    Map<MatchState, Set<MatchState>> jumpMap = characterJumpMap.get(val);
                    if (usePredicateMap) {
                        jumpMap = filterJumpMap(offset, jumpMap, predicateMap);
                    }
                    if (!CollectionUtils.sizeIsEmpty(jumpMap)) {
                        res.add(Pair.<Map<MatchState, Set<MatchState>>, Integer>builder()
                                          .first(jumpMap)
                                          .second(offset + 1)
                                          .build());
                    }
                }

                break;
            }
            default: {
                throw new ParseException("Un-support match input type in matchPf.");
            }
        }
        // 在此处检查所有的 matched state 的 lookAheadExcludes，如果匹配到 LAE，则表示当前没有匹配结果
        List<Pair<Map<MatchState, Set<MatchState>>, Integer>> finalRes = new ArrayList<>(res.size());
        res.forEach(resPair -> {
           Map<MatchState, Set<MatchState>> matchMap = resPair.getFirst();
           Map<MatchState, Set<MatchState>> finalMatchMap = new HashMap<>();
           matchMap.forEach((baseState, directBaseStates) -> {
               Set<MatchState> finalDirectBaseStates = new HashSet<>();
               for (MatchState directBaseState : directBaseStates) {
                   if (!baseState.matchLookaheadExcludes(inputSeq, offset)) {
                       finalDirectBaseStates.add(directBaseState);
                   }
               }
               if (!finalDirectBaseStates.isEmpty()) {
                   finalMatchMap.put(baseState, finalDirectBaseStates);
               }
           });
           if (!finalMatchMap.isEmpty()) {
               finalRes.add(Pair.<Map<MatchState, Set<MatchState>>, Integer>builder()
                                      .first(finalMatchMap)
                                      .second(resPair.getSecond())
                                      .build());
           }
        });
        return finalRes;
    }

    private Map<MatchState, Set<MatchState>> filterJumpMap (int offset,
                                                            Map<MatchState, Set<MatchState>> jumpMap,
                                                            Map<Integer, Set<DefID>> predicateMap) {
        if (CollectionUtils.sizeIsEmpty(jumpMap)) {
            return jumpMap;
        }
        Set<DefID> predicateDefSet = predicateMap.getOrDefault(offset, Collections.emptySet());
        Map<MatchState, Set<MatchState>> filtered = new HashMap<>();
        jumpMap.forEach((base, directStates) -> {
            Set<MatchState> filteredDirectStates = new HashSet<>();
            for (MatchState directState : directStates) {
                if (!directState.isRef() || predicateDefSet.contains(directState.getDefID())) {
                    filteredDirectStates.add(directState);
                }
            }
            if (!filteredDirectStates.isEmpty()) {
                filtered.put(base, filteredDirectStates);
            }
        });
        return filtered;
    }

    private Character tryGetCharVal (MatchUnit inputUnit) {

        if (STRING == inputUnit.getType()) {
            String val = (String) inputUnit.getVal();
            return val.length() == 1 ? val.charAt(0) : null;

        } else {
            return null;
        }
    }

    public boolean matchLookaheadExcludes (List<MatchUnit> inputSeq, int offset) {
        if (lookAheadExcludesMap == null) {
            lookAheadExcludesMap = new StringTreeMap<>();
            if (!CollectionUtils.isEmpty(lookAheadExcludes)) {
                lookAheadExcludes.forEach(lookAheadExcludeItem
                                                  -> lookAheadExcludesMap.put(lookAheadExcludeItem.getVal(), true));
            }
        }
        return !CollectionUtils.isEmpty(lookAheadExcludesMap.matchKey(inputSeq, offset));
    }

    @Builder
    @Getter
    public static class MatchPfRes {

        private Map<MatchState, Set<MatchState>> matchMap;

        private int nextOffset;

    }

    public List<MatchPfRes> matchPfWithRes (List<MatchUnit> inputSeq,
                                            int offset,
                                            Map<Integer, Set<DefID>> predicateMap,
                                            boolean usePredicateMap) {
        List<Pair<Map<MatchState, Set<MatchState>>, Integer>> rawRes
                = matchPf(inputSeq, offset, pfBaseJumpMap, pfStringJumpMap,
                          pfCharacterJumpMap, predicateMap, usePredicateMap);
        return rawRes.stream()
                     .map(resPair -> MatchPfRes.builder()
                                               .matchMap(resPair.getFirst())
                                               .nextOffset(resPair.getSecond())
                                               .build())
                     .collect(Collectors.toList());
    }

    public List<MatchPfRes> matchNextPfWithRes(List<MatchUnit> inputSeq,
                                               int offset,
                                               Map<Integer, Set<DefID>> predicateMap,
                                               boolean usePredicateMap) {

        List<Pair<Map<MatchState, Set<MatchState>>, Integer>> rawRes
                = matchPf(inputSeq, offset, nextBaseJumpMap, nextStringJumpMap,
                          nextCharacterJumpMap, predicateMap, usePredicateMap);
        return rawRes.stream()
                     .map(resPair -> MatchPfRes.builder()
                                               .matchMap(resPair.getFirst())
                                               .nextOffset(resPair.getSecond())
                                               .build())
                     .collect(Collectors.toList());
    }

    public List<Pair<Map<MatchState, Set<MatchState>>, Integer>> matchPf(List<MatchUnit> inputSeq,
                                                                         int offset,
                                                                         Map<Integer, Set<DefID>> predicateMap,
                                                                         boolean usePredicateMap) {
        return matchPf(inputSeq, offset, pfBaseJumpMap, pfStringJumpMap,
                       pfCharacterJumpMap, predicateMap, usePredicateMap);
    }

    public List<Pair<Map<MatchState, Set<MatchState>>, Integer>> matchNextPf(List<MatchUnit> inputSeq,
                                                                             int offset,
                                                                             Map<Integer, Set<DefID>> predicateMap,
                                                                             boolean usePredicateMap) {
        return matchPf(inputSeq, offset, nextBaseJumpMap, nextStringJumpMap,
                       nextCharacterJumpMap, predicateMap, usePredicateMap);
    }

    public boolean isLeafState() {
        return SyntaxLispElementType.isBaseType(src);
    }

    public boolean isInSameDef (MatchState other) {
        return other.getBelongsToDef().getId() == getBelongsToDef().getId();
    }

    public boolean isRef () {
        return SyntaxLispElementType.REF.isTypeMatch(src);
    }

    public String getRefName () {
        return isRef() ? src.getVal() : null;
    }

    public boolean isRepeat () {
        return SyntaxLispElementType.SIMPLIFY_REPEAT.isTypeMatch(src);
    }

    public boolean checkRepeatMatchTimeIsFit (int matchTime) {
        if (!isRepeat()) {
            throw new ParseException("No-repeat can do this func call.");
        }
        return ((SimplifyRepeatElement) src).fit(matchTime);
    }

    public boolean isGreedyMatch () {
        if (!isRepeat()) {
            throw new ParseException("No-repeat can do this func call.");
        }
        return ((SimplifyRepeatElement) src).isGreedyMatch();
    }

    public SyntaxLispElementType getElementType () {
        return src.getType();
    }

    public DefID getBelongsToDef () {
        return belongsToDef;
    }

    public int hashCode () {
        return src.hashCode();
    }

    public String toString () {
        return src.toString();
    }

    public boolean equals (Object o) {
        return (o instanceof MatchState)
                && ((MatchState) o).getSrc().equals(src);
    }


    public boolean isSameRef (String refName) {
        Objects.requireNonNull(refName, "isSameRef param otherOne can not be null.");
        return isRef()
                && src.getVal().equals(refName);
    }

    public int getLayer () {
        return src.getLayer();
    }
}
